from const import *
from runtime import *
import runtime

def inst2str(inst):
    operator, operand = inst
    if operator == OpBinOp:
        operand = BinOps[operand][0]
    elif operator == OpUnOp:
        operand = UnOps[operand][0]
    return "%s  %s" % (VMOps[operator], operand)

class VMException(Exception):pass

class Frame(object):
    def __init__(self, proto, args, upvars, consts, env, upper=None):
        self.proto = proto
        self.insts = proto.insts
        self.args = args
        self.upvars = upvars
        self.consts = consts
        self.env = env
        self.upper = upper

        self.localvars = [None] * proto.maxLocalvars

        if proto.isVararg:
            if len(args) >= proto.argc:
                self.localvars[:proto.argc] = args[:proto.argc]
                self.varargs = runtime.build_list(args[proto.argc:])
            else:
                raise VMException("args:%s argc:%s varg:%s"
                        % (len(args), proto.argc, proto.isVararg))
        else:
            if len(args) == proto.argc:
                self.localvars[:proto.argc] = args[:]
            else:
                raise VMException("args:%s argc:%s varg:%s"
                        % (len(args), proto.argc, proto.isVararg))

        self.stack = []

    def buildFromClosure(self, closure, args):
        return Frame(closure.proto, args, closure.upvars, self.consts,
                self.env, self)

    def eval(self):
        pc = 0
        insts = self.insts

        while True:
            try:
                inst = insts[pc]
            except IndexError:
                return 
            operator, operand = inst
            # print inst2str(inst)

            if operator == OpRet:
                self.closeUpvars(0)

                return self.stack[-1*operand:]

            elif operator == OpLoadLocal:  
                assert operand >= 0
                var = self.localvars[operand]
                assert var is not None
                self.stack.append(var)

            elif operator == OpSetLocal:
                assert operand >= 0
                a = self.stack.pop()
                self.localvars[operand] = a

            elif operator == OpLoadGlobal:
                a = self.consts[operand]
                self.stack.append(self.env[a])

            elif operator == OpSetGlobal:
                b = self.stack.pop()
                s = self.consts[operand]
                self.env[s] = b

            elif operator == OpLoadUpvar:
                assert operand >= 0 
                var = self.upvars[operand].get()
                self.stack.append(var)

            elif operator == OpSetUpvar:
                a = self.stack.pop()
                assert operand >= 0 
                self.upvars[operand].set(a)

            elif operator == OpLoadVarg:
                self.stack.append(self.varargs)

            elif operator ==  OpLoadConst:
                try:
                    var = self.consts[operand]
                    self.stack.append(var)
                except IndexError:
                    raise VMException("Error const index '%s' " % operand)

            elif operator == OpBinOp:
                b = self.stack.pop()
                a = self.stack.pop()
                try: 
                    self.stack.append(BinOps[operand][1](a,b))
                except KeyError:
                    raise VMException("Unknown BinOp '%s' " % operand)

            elif operator == OpUnOp:
                a = self.stack.pop()
                try:
                    symbol, unop = UnOps[operand]
                    self.stack.append(unop(a))
                except KeyError:
                    raise VMException("Unknown UnOp '%s' " % operand)
                    
            elif operator == OpCall:
                argc = operand
                closure = self.stack[-1 * (1 + argc)]
                if argc == 0:
                    args = []
                else:
                    args = self.stack[-1 * argc :]
                for i in range(0, argc + 1):
                    self.stack.pop()

                if isinstance(closure, runtime.Closure):
                    frame = self.buildFromClosure(closure, args)
                    rets = frame.eval()
                else:
                    rets = closure(*args)
                self.stack.extend(rets)

            elif operator == OpTailCall:
                argc = operand
                closure = self.stack[-1 * (1 + argc)]
                if argc == 0:
                    args = []
                else:
                    args = self.stack[-1 * argc :]

                if isinstance(closure, runtime.Closure):
                    self.__init__(proto=closure.proto, args=args, 
                            upvars=closure.upvars, consts=self.consts,
                            env=self.env, upper=self.upper
                            )
                    pc = 0
                    insts = self.insts
                    continue
                else:
                    rets = closure(*args)
                    self.stack.extend(rets)
 
            elif operator == OpJump:
                pc += operand
                continue

            elif operator == OpTest:
                """
                (if (= x 2) (a) (b)) 
                # In the end , only the value of whole "if" 
                # expression is on the top of the stack.
                Load      x
                LoadConst 2
                EQ
                TEST 3
                LOAD A
                JUMP 2
                LOAD B
                """
                isTrue = self.stack.pop()
                if not isTrue:
                    pc += operand
                    continue

            elif operator == OpPop:
                for i in range(0, operand):
                    self.stack.pop()

            elif operator == OpCloseUpvar:
                # FIXME: not efficient way
                for upvar in reversed(self.upvars):
                    if isinstance(upvar, runtime.UpvarForVarg):
                        upvar.close()
                    elif (upvars.upvar.stat == Upvar.OPEN
                        and upvar.index >= operand):
                        upvar.close()

            elif operator == OpBuildClosure:
                a = self.stack.pop()
                closure = Closure(a, self)
                self.stack.append(closure)

            else:
                raise VMException("Unknown VM operator with operand %s:%s ." 
                        % (operator,operand))
            pc += 1

    def closeUpvars(self, index):
        for upvar in self.upvars:
            if isinstance(upvar, runtime.UpvarForVarg):
                upvar.close()
            elif upvar.index >= index and upvar.stat == upvar.OPEN:
                upvar.close()

class VM(object):
    def __init__(self):
        self.env = {}

    def start(self):
        entry = self.consts[-1]
        main = Frame(proto=entry, args=[], 
                upvars=[], consts = self.consts, env=self.env )
        main.eval()

if __name__ == '__main__':
    pass
